# Load necessary packages
library(tidyverse)
library(nnet)

# Load data
df <- readRDS("kwit_tb_limited.rds")

# remove persons with missing HIV 
df.hiv <- df %>% filter(included_hiv_narm)

# Model fit
fit1 <- multinom(treatment_outcome_new ~ hiv_art_cat + 
                   as.factor(age_cat4) + 
                   index_gender + 
                   tb_type + 
                   tobacco_use + 
                   alcohol_use, data=df.hiv)

# Coefficient estimates
est <- round(t(exp(summary(fit1)$coefficients)),2)[-1,]

# Standard error (Delta method)
var <- (t(summary(fit1)$standard.errors)^2)[-1,]
se <- round(sqrt((est^2)*var),2)

# 95% confidence intervals 
ci_low <- round(t(exp(summary(fit1)$coefficients + qnorm(.025)*summary(fit1)$standard.errors)),2)[-1,]
ci_up <- round(t(exp(summary(fit1)$coefficients + qnorm(.975)*summary(fit1)$standard.errors)),2)[-1,]

# P-values
pvals <- round(t(pnorm(abs(summary(fit1)$coefficients/summary(fit1)$standard.errors),lower.tail = FALSE)),3)[-1,]

# Make separate data frames for results
## Down referred vs. Unfavorable 
var.names <- c("HIV+, ART (vs. HIV+, no ART)","HIV- (vs. HIV+, no ART)",
               "Age: 25-54 (vs. 15-24)",
               "Age: 55-64 (vs. 15-24)",
               "Age: 65+ (vs. 15-24)",
               "Gender: Male (vs. Female)",
               "TB Type: Extra Pulmonary (vs. Pulmonary)",
               "TB Type: Both (vs. Pulmonary)",
               "TB Type: Unanswered (vs. Pulmonary)",
               "Tobacco: Unanswered (vs. No)",
               "Tobacco: Yes (vs. No)",
               "Alcohol: Unanswered (vs. No)",
               "Alcohol: Yes (vs.No)")
col.names <- c("Odds Ratio","Standard error","95% CI","P-value")

df1 <- data.frame(or=est[,1],
                  se=se[,1],
                  ci=paste0("(",ci_low[,1],", ",ci_up[,1],")"),
                  pval=pvals[,1])
rownames(df1) <- var.names
colnames(df1) <- col.names

## Success vs. Unfavorable 
df2 <- data.frame(or=est[,2],
                  se=se[,2],
                  ci=paste0("(",ci_low[,2],", ",ci_up[,2],")"),
                  pvals[,2])
rownames(df2) <- var.names
colnames(df2) <- col.names

# Save results
saveRDS(df1,"model_results_downreferred.rds")
saveRDS(df2,"model_results_success.rds")
